# -*- coding: utf-8 -*-


import os
import numpy as np
from netCDF4 import Dataset

from tkits.dt import dt_range
from tkits.Path import list_dir, joinext, mkdirs
from metlib.data.basket import loadbasket
from fg_solar_sat_common import sat_lonlat_range


solar_sat_dir = '/home/yh/work/dataset/solar/cma'
stored_dir = '/home/yh/solar'


def create_solar_sat_nc(fname, jy, ix, tag, stat_type='PT', **kwargs):
    solar = Dataset(fname, 'w')
    # dimensions
    solar.createDimension('jy', size=jy)
    solar.createDimension('ix', size=ix)
    if stat_type == 'PT':
        solar.createDimension('date', size=tag)
        solar.createVariable('date', str, dimensions=('date',))
        solar.createVariable('rad', 'f4', dimensions=('date', 'jy', 'ix'))
    elif stat_type in ['PM', 'PY']:
        solar.createDimension('tag', size=tag)
        solar.createVariable('tag', str, dimensions=('tag',))
        solar.createVariable('mean', 'f4', dimensions=('tag', 'jy', 'ix'))
        solar.createVariable('std', 'f4', dimensions=('tag', 'jy', 'ix'))
        solar.createVariable('sum', 'f4', dimensions=('tag', 'jy', 'ix'))
        solar.createVariable('count', 'i2', dimensions=('tag',))
    # variables
    solar.createVariable('lat', 'f4', dimensions=('jy', 'ix'))
    solar.createVariable('lon', 'f4', dimensions=('jy', 'ix'))
    # attributes
    solar.max_lon = kwargs.get('max_lon')
    solar.min_lon = kwargs.get('min_lon')
    solar.max_lat = kwargs.get('max_lat')
    solar.min_lat = kwargs.get('min_lat')
    solar.lon_delta = kwargs.get('lon_delta')
    solar.lat_delta = kwargs.get('lat_delta')
    solar.sat_name = kwargs.get('sat')
    return solar


def npy2nc(sat='fy2c'):
    fy_sat_dir = os.path.join(solar_sat_dir, sat)

    lonlat = {}
    loadbasket(os.path.join(fy_sat_dir, 'lonlat.zip'), dest=lonlat)
    lon = lonlat['LON']
    lat = lonlat['LAT']
    jy, ix = lon.shape

    # nan array
    nan_data = np.empty((jy, ix), dtype='f4')
    nan_data[:] = np.nan

    for year in sat_lonlat_range[sat]['years']:
        begin_dt = '%s0101' % year
        end_dt = '%s0101' % (int(year) + 1)
        dt_list = dt_range(begin_dt, end_dt, step='1d')
        npy_names = sorted(list_dir(fy_sat_dir, pattern=year, path=False))
        # create PT nc
        fy_stored_dir = os.path.join(stored_dir, sat, 'PT', year)
        mkdirs(fy_stored_dir)
        nc_name = os.path.join(fy_stored_dir, joinext('rad', 'nc'))
        solar = create_solar_sat_nc(nc_name, jy, ix, len(dt_list),
                                    stat_type='PT', **sat_lonlat_range[sat])
        solar.variables['lon'][:] = lon
        solar.variables['lat'][:] = lat

        date = solar.variables['date']
        for index, dt in enumerate(dt_list):
            dt_str = dt.strftime('%Y%m%d')
            npy_name = joinext(dt_str, 'npy')
            if npy_name in npy_names:
                npy_path = os.path.join(fy_sat_dir, npy_name)
                solar.variables['rad'][index, :] = np.load(npy_path)
            else:
                solar.variables['rad'][index, :] = nan_data
            date[index] = dt_str
        solar.close()


def fy_stat_year(sat):
    years = sat_lonlat_range[sat]['years']
    for year in years:
        sat_ts = Dataset(os.path.join(stored_dir, sat, 'PT', year, 'rad.nc'))
        lon = sat_ts.variables['lon']
        lat = sat_ts.variables['lat']
        rad = sat_ts.variables['rad']
        jy, ix = lon.shape

        py_dir = os.path.join(stored_dir, sat, 'PY', year)
        mkdirs(py_dir)
        nc_name = os.path.join(py_dir, 'rad.nc')
        stat_year = create_solar_sat_nc(nc_name, jy, ix, 1, stat_type='PY', **sat_lonlat_range[sat])
        stat_year.variables['lon'][:] = lon
        stat_year.variables['lat'][:] = lat
        stat_year.variables['tag'][0] = year
        stat_year.variables['mean'][0,:] = np.nanmean(rad, axis=0)
        stat_year.variables['count'][0] = rad[:,0,0][np.isfinite(rad[:,0,0])].size
        stat_year.variables['std'][0,:] = np.nanstd(rad, axis=0)
        stat_year.variables['sum'][0,:] = np.nansum(rad, axis=0)
        stat_year.close()
        sat_ts.close()


def fy_stat_month(sat):
    years = sat_lonlat_range[sat]['years']
    for year in years:
        sat_ts = Dataset(os.path.join(stored_dir, sat, 'PT', year, 'rad.nc'))
        lon = sat_ts.variables['lon']
        lat = sat_ts.variables['lat']
        rad = sat_ts.variables['rad']
        jy, ix = lon.shape

        pm_dir = os.path.join(stored_dir, sat, 'PM', year)
        mkdirs(pm_dir)
        nc_name = os.path.join(pm_dir, 'rad.nc')
        stat_month = create_solar_sat_nc(nc_name, jy, ix, 12, stat_type='PM', **sat_lonlat_range[sat])
        stat_month.variables['lon'][:] = lon
        stat_month.variables['lat'][:] = lat
        fmt = '%s%02d'
        # 12 months
        date = sat_ts.variables['date'][:]
        for i in range(12):
            # only str
            current_month = fmt % (year, i+1)
            next_month = fmt % (year, i+2)
            stat_month.variables['tag'][i] = current_month
            a = (date > current_month)
            b = (date < next_month)
            rad_month = rad[a & b]
            # month mean
            stat_month.variables['count'][i] = rad_month[:,0,0][np.isfinite(rad_month[:,0,0])].size
            stat_month.variables['mean'][i,:] = np.nanmean(rad_month, axis=0)
            stat_month.variables['std'][i,:] = np.nanstd(rad_month, axis=0)
            stat_month.variables['sum'][i,:] = np.nansum(rad_month, axis=0)
        stat_month.close()
        sat_ts.close()


if __name__ == '__main__':
    for sat in ['fy2c', 'fy2d', 'fy2e', 'fy2f']:
        # print 'pt'
        # npy2nc(sat)
        print 'py'
        fy_stat_year(sat)
        print 'pm'
        fy_stat_month(sat)
